#该问题发生在我要将图片和其对应的xml文件划分成相对应的测试集和训练集
#文件夹排版
# ├─Annotations
    #标注的xml文件
# └─images
    #图片
#其实就是生成3个txt文件   在--txt-save-path目录下生成 trainval.txt  val.txt  train.txt
#生成后目录结构如下
# ├── Annotations
    #标注的xml文件
# ├── images
    #图片
# └── ImagesSets
#     └── Main
         #train.txt、trainval.txt、val.txt

#接下去进入xml_to_yolo.py进行后续处理
import argparse
from pathlib import Path
import random

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="img、xmlToValTrain")
    parser.add_argument('--root-dir', type=str, default="/home/luoluoluo/data/dataset/elevator/train_detect")
    return parser.parse_args()

def main():
    args = Parse_Arguments()
    trainval_percent = 1.0  # 训练集和验证集所占比例。 这里没有划分测试集
    train_percent = 0.8    # 训练集所占比例，可自己进行调整

    #数据集的划分，地址选择自己数据下的ImageSets/Main
    txt_save_path = args.root_dir + '/ImagesSets/Main'
    #xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
    xml_file_path = args.root_dir + '/Annotations'
    total_xml = Path(xml_file_path).glob("*.xml")
    total_xml = [f for f in total_xml]
    num = len(total_xml)
    if not Path(txt_save_path).exists():
        Path(txt_save_path).mkdir(parents=True)
    list_index = range(num)
    train_val_num = int(num * trainval_percent)
    train_num = int(train_val_num * train_percent)
    train_val_index = random.sample(list_index, train_val_num)
    train_index = random.sample(train_val_index, train_num)

    file_train_val = open(txt_save_path + '/trainval.txt', 'w')
    file_test = open(txt_save_path + '/val.txt', 'w')
    file_train = open(txt_save_path + '/train.txt', 'w')
    file_val = open(txt_save_path + '/val.txt', 'w')

    for i in list_index:
        name = total_xml[i].stem + '\n'
        if i in train_val_index:
            file_train_val.write(name)
            if i in train_index:
                file_train.write(name)
            else:
                file_val.write(name)
        else:
            file_test.write(name)
    print(f"创建trainval、val、train txt文件成功")
    file_train_val.close()
    file_test.close()
    file_train.close()
    file_val.close()

if __name__ == "__main__":
    main()
